#include "GeneratePartition.h"

#include <algorithm>
#include <iostream>
#include <vector>

#include "gtest/gtest.h"
#include "metis.h"
#include <gtest/gtest.h>

#include "Utility/Algebra/Matrix.h"
#include "Utility/Configuration/SGParameter.h"

#include "DBManager/DBServiceFactory.h"
#include "DBManager/IService/Input/Element/IElementService.h"

#include "SGFem/Pre/DOFHandler/MinimumDegreeSort.h"

#include "SGFem/ElementCalculator/Factory/ElementCalculatorFactory.h"

using SG::Algebra::Graph;
using SG::Algebra::Matrix;

using SG::DataStructure::Common::Id;
using SG::DBManager::DBServiceFactorySP;
using SG::FEMSolver::Element::ElementCalculatorFactory;
using namespace SG::FEMSolver;
using namespace SG::DBManager;
using IdxSet = std::vector<idx_t>;

GeneratePartitioning::GeneratePartitioning ()
{}

void GeneratePartitioning::SGPartitioner (SG::DBManager::DBServiceFactorySP  dbServiceFactory,
                                          SG::DataStructure::Common::IdArray NodeIds,
                                          SG::DataStructure::Common::IdArray ElementIds,
                                          SG::Algebra::Matrix                ElementNodeIds,
                                          const SG::Utility::SGParameter&    SGParameters,
                                          _OUT SG::DataStructure::Common::IdArray& SGPartitionElement)
{
#ifdef SG_USE_METIS
    // 将生成的图传入metis接口， 得到新的排序
    auto g = GeneratePartitioning::GenerateElementGraph (dbServiceFactory, NodeIds, ElementIds, ElementNodeIds);
    // 调用Metis图分割程序
    GeneratePartitioning::GeneratePartition (ElementIds, g, SGParameters, SGPartitionElement);
    // 输出网格分区及索引
    SG::DataStructure::Common::IdArray PartitionIndex;
    GeneratePartitioning::GenerateGlobalIndex (dbServiceFactory, SGPartitionElement, SGParameters, NodeIds, ElementIds, PartitionIndex);
#else
    std::cout << "Mesh::GeneratePartitioning(...): "
                 "MFEM was compiled without Metis."
              << std::endl;
    // return NULL;

#endif
}

Graph GeneratePartitioning::GenerateElementGraph (const SG::DBManager::DBServiceFactorySP     dbServiceFactory,
                                                  std::vector<SG::DataStructure::Common::Id>& nodeIds,
                                                  const SG::DataStructure::Common::IdArray    ElementIds,
                                                  SG::Algebra::Matrix                         ElementNodeIds)
{
    const auto& eleSumInfoService = dbServiceFactory->get<IElementService> ();

    // 第一次遍历拓扑关系，统计每个单元对应连接单元关系。
    std::size_t              ElementSize = ElementIds.size ();
    std::size_t              nodeSize    = nodeIds.size ();
    std::vector<std::size_t> nEdge (ElementSize, 0);
    auto                     eleTypes = eleSumInfoService->GetAllElementTypes ();
    std::size_t              posi;
    std::vector<std::size_t> EleIndexes (ElementSize);

    for (std::size_t inode = 0; inode < ElementSize; ++inode)
    {
        EleIndexes[inode] = inode;
    }
    std::vector<std::size_t> NodeIndexes (nodeSize);
    for (std::size_t inode = 0; inode < nodeSize; ++inode)
    {
        NodeIndexes[inode] = nodeIds[inode];
    }

    for (const auto& eleType : eleTypes)
    {
        // 跳过质量单元
        //  if (/* condition */)
        //  {
        //      /* code */
        //  }
        Matrix eleNodes = eleSumInfoService->GetElementNodes (eleType);  // 需要建立类型和数据表的关联。
        auto   nEleNode = eleNodes.GetColSize ();
        // 单元没有节点或者仅有一个节点
        if (nEleNode < 2)
            continue;
        // 对节点循环
        std::vector<SG::DataStructure::Common::Id> NeighVecAll (1);
        NeighVecAll[0] = -1;
        for (std::size_t inode = 0; inode < nodeSize; ++inode)
        {
            // auto                                       it = AlleleNode.begin ();
            std::vector<SG::DataStructure::Common::Id> Neigh;
            for (std::size_t iEle = 0; iEle < ElementSize; ++iEle)
            {
                for (std::size_t jEle = 0; jEle < ElementNodeIds.GetColSize (); ++jEle)
                {
                    if (ElementNodeIds[iEle][jEle] == NodeIndexes[inode])
                    {
                        Neigh.push_back (iEle);
                        break;
                    }
                }
            }
            if (Neigh.size () == 1)
                continue;
            std::vector<SG::DataStructure::Common::Id> Neigh1;
            Neigh1.insert (Neigh1.end (), Neigh.begin (), Neigh.end ());
            Neigh1.push_back (-1);

            auto it = std::search (NeighVecAll.begin (), NeighVecAll.end (), Neigh1.begin (), Neigh1.end ());

            if (it != NeighVecAll.end ())
            {
                // std::cout << "子向量存在于主向量中" << std::endl;
                continue;
            }
            else
            {
                // std::cout << "子向量不存在于主向量中" << std::endl;
                NeighVecAll.insert (NeighVecAll.end (), Neigh1.begin (), Neigh1.end ());
            }

            for (std::size_t i = 0; i < Neigh.size (); ++i)
            {
                for (std::size_t j = i + 1; j < Neigh.size (); ++j)
                {
                    ++(nEdge[Neigh[i]]);
                    ++(nEdge[Neigh[j]]);
                }
            }
        }
    }
    if (1)
    {
        std::size_t nodeIndex = 0;
        for (auto vSize : nEdge)
        {
            if (vSize < 2)
            {
                std::cout << "independented node " << nodeIds[nodeIndex] << std::endl;
            }
            ++nodeIndex;
        }
    }
    // 生成图的结构
    Graph g (nEdge);
    // 再次遍历单元，完成对各个节点邻居的记录
    // pIndex, 用于填入边位置进行计数
    std::vector<std::size_t> pIndex = g.m_pointer;

    for (const auto& eleType : eleTypes)
    {
        // 跳过质量单元
        //  if (/* condition */)
        //  {
        //      /* code */
        //  }
        Matrix eleNodes = eleSumInfoService->GetElementNodes (eleType);  // 需要建立类型和数据表的关联。
        auto   nEleNode = eleNodes.GetColSize ();
        // 单元没有节点或者仅有一个节点
        if (nEleNode < 2)
            continue;
        // 对节点循环
        std::vector<SG::DataStructure::Common::Id> NeighVecAll (1);
        NeighVecAll[0] = -1;
        for (std::size_t inode = 0; inode < nodeSize; ++inode)
        {
            std::vector<SG::DataStructure::Common::Id> Neigh;
            for (std::size_t iEle = 0; iEle < ElementSize; ++iEle)
            {
                for (std::size_t jEle = 0; jEle < ElementNodeIds.GetColSize (); ++jEle)
                {
                    if (ElementNodeIds[iEle][jEle] == NodeIndexes[inode])
                    {
                        Neigh.push_back (iEle);
                        break;
                    }
                }
            }

            if (Neigh.size () == 1)
                continue;
            std::vector<SG::DataStructure::Common::Id> Neigh1;

            Neigh1.insert (Neigh1.end (), Neigh.begin (), Neigh.end ());
            Neigh1.push_back (-1);

            auto it = std::search (NeighVecAll.begin (), NeighVecAll.end (), Neigh1.begin (), Neigh1.end ());

            if (it != NeighVecAll.end ())
            {
                // std::cout << "子向量存在于主向量中" << std::endl;
                continue;
            }
            else
            {
                // std::cout << "子向量不存在于主向量中" << std::endl;
                NeighVecAll.insert (NeighVecAll.end (), Neigh1.begin (), Neigh1.end ());
            }

            for (std::size_t i = 0; i < Neigh.size (); ++i)
            {
                for (std::size_t j = i + 1; j < Neigh.size (); ++j)
                {
                    posi                = pIndex[Neigh[i]];
                    g.m_neighbour[posi] = EleIndexes[Neigh[j]];
                    ++pIndex[Neigh[i]];

                    posi                = pIndex[Neigh[j]];
                    g.m_neighbour[posi] = Neigh[i];
                    ++pIndex[Neigh[j]];
                }
            }
        }
    }

    // 消除重复边记录。 一条图边可能被两个及以上单元共用。存在重复记录
    int mpt = g.m_pointer[0];
    for (std::size_t vI = 0; vI < g.m_pointer.size () - 1; ++vI)
    {
        // 将顶点 vI 的邻居copy 到 nerghborI
        std::vector<int> nerghborI;
        nerghborI.assign (g.m_neighbour.begin () + mpt, g.m_neighbour.begin () + g.m_pointer[vI + 1]);
        // 删除重复出现的邻居
        std::sort (nerghborI.begin (), nerghborI.end ());
        nerghborI.erase (unique (nerghborI.begin (), nerghborI.end ()), nerghborI.end ());
        auto len = nerghborI.size ();

        // 将邻居顶点放在正确的位置
        const auto& posi = g.m_pointer[vI];
        for (std::size_t i = 0; i < len; ++i)
        {
            g.m_neighbour[i + posi] = nerghborI[i];
        }

        mpt = g.m_pointer[vI + 1];  // 记录下一个顶点邻居的起始位置
        // 更新 vI 邻居的结束位置
        g.m_pointer[vI + 1] = g.m_pointer[vI] + len;
    }
    // 删除 neighbour中无效数据
    g.m_neighbour.resize (*(g.m_pointer.end () - 1));
    return g;
}

void GeneratePartitioning::GeneratePartition (const SG::DataStructure::Common::IdArray ElementIds,
                                              SG::Algebra::Graph                       g,
                                              const SG::Utility::SGParameter&          SGParameters,
                                              _OUT SG::DataStructure::Common::IdArray& SGPartitionElement)
{
    // If running in parallel, print messages only from rank 0.
    //  #ifdef MFEM_USE_MPI
    //     int init_flag, fin_flag;
    //     MPI_Initialized(&init_flag);
    //     MPI_Finalized(&fin_flag);
    //     if (init_flag && !fin_flag)
    //     {
    //        int rank;
    //        MPI_Comm_rank(GetGlobalMPI_Comm(), &rank);
    //        if (rank != 0) { print_messages = 0; }
    //     }
    //  #endif

    std::size_t              nelements = ElementIds.size ();
    std::vector<std::size_t> indexes;
    std::size_t              nparts{0};
    std::string              part_method;

    if (SGParameters.Has ("NParts"))
    {
        nparts = SGParameters.GetValue ("NParts").GetInt ();
    }
    // 1. 单分区直接返回0
    if (nparts == 1)
    {
        SG::DataStructure::Common::IdArray SGPartitionElementtmp (nelements);
        for (std::size_t i = 0; i < nelements; ++i)
        {
            SGPartitionElementtmp[i] = 0;
        }
        SGPartitionElement = SGPartitionElementtmp;
    }

    // 2. 元素数<=分区数时每个元素单独分区
    if (nelements <= nparts)
    {
        SG::DataStructure::Common::IdArray SGPartitionElementtmp (nelements);
        for (std::size_t i = 0; i < nelements; ++i)
        {
            SGPartitionElementtmp[i] = i;
        }
        SGPartitionElement = SGPartitionElementtmp;
    }

    // 3. 否则直接调用Metis实现分区
    if (SGParameters.Has ("Method"))
    {
        part_method = SGParameters.GetValue ("Method").GetString ();
    }

    idx_t  nvtxs = g.GetNodeSize ();
    idx_t  ncon  = 1;
    IdxSet xadj (nvtxs + 1);
    IdxSet adjncy (g.m_neighbour.size ());
    for (std::size_t i = 0; i < g.m_pointer.size () + 1; ++i)
    {
        xadj[i] = g.m_pointer[i];
    }
    for (std::size_t i = 0; i < g.m_neighbour.size (); ++i)
    {
        adjncy[i] = g.m_neighbour[i];
    }
    // idx_t wgtflag = 0;
    // idx_t numflag = 0;
    idx_t mparts = ( std::size_t ) nparts;
    idx_t options[40];
    METIS_SetDefaultOptions (options);
    options[METIS_OPTION_CONTIG] = 1;
    idx_t edgecut;
    idx_t errflag;
    // In case METIS have been compiled with 64bit indices
    // bool freedata = false;
    IdxSet mpartitioning (nvtxs, 0);

    // This function should be used to partition a graph into a small
    // number of partitions (less than 8).

    if (part_method == "METIS_PartGraphRecursive")
    {
        errflag = METIS_PartGraphRecursive (
            &nvtxs, &ncon, xadj.data (), adjncy.data (), NULL, NULL, NULL, &mparts, NULL, NULL, options, &edgecut, mpartitioning.data ());
        if (errflag != 1)
        {
            std::cout << "Mesh::GeneratePartitioning: " << std::endl << " error in METIS_PartGraphRecursive!" << std::endl;
        }
    }
    // This function should be used to partition a graph into a large
    // number of partitions (greater than 8).
    if (part_method == "METIS_PartGraphKway")
    {
        errflag = METIS_PartGraphKway (
            &nvtxs, &ncon, xadj.data (), adjncy.data (), NULL, NULL, NULL, &mparts, NULL, NULL, options, &edgecut, mpartitioning.data ());
        if (errflag != 1)
        {
            std::cout << "Mesh::GeneratePartitioning: " << std::endl << " error in METIS_PartGraphRecursive!" << std::endl;
        }
    }

    // The objective of this partitioning is to minimize the total
    // communication volume
    // if (part_method == 2 || part_method == 5)
    // {
    //    METIS_PartGraphVKway(&nvtxs,
    //                         &ncon,
    //                         xadj.data(),
    //                         adjncy.data(),
    //                         NULL,
    //                         NULL,
    //                         NULL,
    //                         &mparts,
    //                         NULL,
    //                         NULL,
    //                         options,
    //                         &edgecut,
    //                         mpartitioning);
    // }
    SG::DataStructure::Common::IdArray SGPartitionElementtmp (nelements);
    std::vector<int> parts;
    parts.assign (mpartitioning.begin (), mpartitioning.end ());
    for (std::size_t i = 0; i < parts.size (); ++i)
    {
        SGPartitionElementtmp[i] = parts[i];
    }
    SGPartitionElement = SGPartitionElementtmp;
}

void GeneratePartitioning::GenerateGlobalIndex (const SG::DBManager::DBServiceFactorySP dbServiceFactory,
                                                SG::DataStructure::Common::IdArray&     SGPartitionElement,
                                                const SG::Utility::SGParameter&         SGParameters,
                                                SG::DataStructure::Common::IdArray      NodeIds,
                                                SG::DataStructure::Common::IdArray      ElementIds,
                                                SG::DataStructure::Common::IdArray _OUT PartitionIndex)
{}
